import time
from ultralytics import YOLO

def train():
    # yolov8n模型训练：训练模型的数据为'A_my_data.yaml'，轮数为100，图片大小为640，设备为本地的GPU显卡，关闭多线程的加载，图像加载的批次大小为4，开启图片缓存
    model = YOLO('model/yolov8n.pt')  # load a pretrained model (recommended for training)
    results = model.train(data='../../ultralytics/datasets/Helmet.yaml', epochs=10, imgsz=640, device=[], workers=0, batch=1, cache=True)  # 开始训练
    time.sleep(10) # 睡眠10s，主要是用于服务器多次训练的过程中使用


#测试
def val():
    # 加载自己训练好的模型，填写相对于这个脚本的相对路径或者填写绝对路径均可
    model = YOLO("runs/detect/train/weights/best.pt")

    # 开始进行验证，验证的数据集为'A_my_data.yaml'，图像大小为640，批次大小为4，置信度分数为0.25，交并比的阈值为0.6，设备为0，关闭多线程（windows下使用多线程加载数据容易出现问题）
    validation_results = model.val(data='../../ultralytics/datasets/Helmet.yaml', imgsz=640, batch=4, conf=0.25, iou=0.6, device=[], workers=0)
    print(validation_results)

def predict():
    # Load a model
    model = YOLO("runs/detect/train2/weights/best.pt")  # pretrained YOLOv8n model

    # Run batched inference on a list of images
    results = model(["image/demo.png", ])  # return a list of Results objects

    print(results)

    # Process results list
    for result in results:
        boxes = result.boxes  # Boxes object for bounding box outputs
        masks = result.masks  # Masks object for segmentation masks outputs
        keypoints = result.keypoints  # Keypoints object for pose outputs
        probs = result.probs  # Probs object for classification outputs
        obb = result.obb  # Oriented boxes object for OBB outputs
        result.show()  # display to screen
        result.save(filename="image/demo_result.jpg")  # save to disk


if __name__ == '__main__':
    # train()
    # val()
    predict()



